package com.example.demo.demos.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author lyh
 * @desc 处理WebSocket请求的接口类, 移植到需要登录验证的项目时需要添加WebSocket请求至白名单(非常重要)
 * @date 2025/7/21 16:57
 */
public class AuthWebSocketHandler extends TextWebSocketHandler {

    // 存储所有活动WebSocket会话 (sessionId -> session)
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final ScheduledExecutorService heartExecutor = Executors.newSingleThreadScheduledExecutor();

    // 初始化心跳任务（5秒一次）
    public AuthWebSocketHandler() {
        heartExecutor.scheduleAtFixedRate(this::sendHeartbeat, 5, 5, TimeUnit.SECONDS);
    }

    /**
     * 当新连接建立时触发
     *
     * @param session 新建立的WebSocket会话
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.put(session.getId(), session);
        // 添加连接日志
        try {
            InetSocketAddress remoteAddress = (InetSocketAddress) session.getRemoteAddress();
            System.out.println("新客户端连接: " + remoteAddress.getAddress().getHostAddress());
        } catch (Exception e) {
            System.out.println("新客户端连接（无法获取IP）");
        }
    }

    /**
     * 处理来自客户端的文本消息
     *
     * @param session 发送消息的会话
     * @param message 接收到的消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        Map<String, String> data = objectMapper.readValue(message.getPayload(), Map.class);
        String type = data.get("type");

        if ("C_TO_S_AUTH_ACTIVE".equals(type)) {
            // todo 认证逻辑处理
            boolean isValid = validateCode(data.get("code"));

            // 返回认证结果
            Map<String, Object> response = new HashMap<>();
            response.put("type", "C_TO_S_AUTH_ACTIVE");
            response.put("result", isValid);
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
        }
    }

    /**
     * 获取所有已连接客户端的IP地址列表
     *
     * @return 客户端IP地址列表（可能包含重复）
     */
    public List<String> getAllClientIps() {
        return sessions.values().stream()
                .filter(WebSocketSession::isOpen)
                .map(session -> {
                    try {
                        InetSocketAddress socketAddress = (InetSocketAddress) session.getRemoteAddress();
                        InetAddress address = socketAddress.getAddress();

                        // 处理IPv6环回地址的情况
                        if (address instanceof Inet6Address && address.isLoopbackAddress()) {
                            return "127.0.0.1"; // 显式返回IPv4环回地址
                        }

                        // 处理IPv4映射的IPv6地址 (::ffff:192.168.x.x)
                        if (address instanceof Inet6Address) {
                            byte[] bytes = address.getAddress();
                            if (bytes.length == 16 &&
                                    bytes[0] == 0 && bytes[1] == 0 &&
                                    bytes[2] == 0 && bytes[3] == 0 &&
                                    bytes[4] == 0 && bytes[5] == 0 &&
                                    bytes[6] == 0 && bytes[7] == 0 &&
                                    bytes[8] == 0 && bytes[9] == 0 &&
                                    bytes[10] == (byte) 0xff && bytes[11] == (byte) 0xff) {

                                // 提取真实的IPv4地址
                                return String.format("%d.%d.%d.%d",
                                        bytes[12] & 0xff,
                                        bytes[13] & 0xff,
                                        bytes[14] & 0xff,
                                        bytes[15] & 0xff);
                            }
                        }

                        // 标准IPv4地址
                        return address.getHostAddress();
                    } catch (Exception e) {
                        return "未知IP";
                    }
                })
                .collect(Collectors.toList());
    }

    /**
     * 主动断开指定会话连接
     *
     * @param sessionId 要断开的会话ID
     */
    public void closeSession(String sessionId) throws Exception {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            session.close(CloseStatus.NORMAL);
            sessions.remove(sessionId);
        }
    }

    /**
     * 向所有连接的客户端广播消息
     *
     * @param message 要发送的消息内容
     */
    public void broadcastMessage(String message) {
        sessions.values().forEach(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(message));
                }
            } catch (IOException e) {
                // 记录发送失败日志
                System.err.println("消息发送失败至会话: " + session.getId() + " 错误: " + e.getMessage());
            }
        });
    }

    /**
     * 向指定会话发送消息
     *
     * @param sessionId 目标会话ID
     * @param message   要发送的消息内容
     * @return 发送成功返回true，失败返回false
     */
    public boolean sendToSession(String sessionId, String message) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
                return true;
            } catch (IOException e) {
                System.err.println("定向消息发送失败: " + sessionId + " 错误: " + e.getMessage());
            }
        }
        return false;
    }

    // 发送心跳消息
    private void sendHeartbeat() {
        Map<String, String> heartMsg = new HashMap<>();
        heartMsg.put("type", "S_TO_C_HEART");
        try {
            String jsonMsg = objectMapper.writeValueAsString(heartMsg);
            sessions.values().forEach(session -> {
                try {
                    if (session.isOpen()) {
                        session.sendMessage(new TextMessage(jsonMsg));
                    }
                } catch (IOException e) {
                    System.err.println("心跳发送失败: " + session.getId());
                }
            });
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private boolean validateCode(String code) {
        // 实际业务逻辑（如数据库校验）
        return code != null && code.startsWith("saga");
    }

    /**
     * 当连接关闭时触发
     *
     * @param session 关闭的会话
     * @param status  关闭状态
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        sessions.remove(session.getId());
    }

    /**
     * 向指定IP列表发送消息
     *
     * @param targetIps 目标IP列表
     * @param message   要发送的消息
     */
    public void sendToIps(List<String> targetIps, String message) {
        sessions.values().stream()
                .filter(WebSocketSession::isOpen)
                .filter(session -> {
                    String sessionIp = getIpFromSession(session);
                    return targetIps.contains(sessionIp);
                })
                .forEach(session -> {
                    try {
                        session.sendMessage(new TextMessage(message));
                    } catch (IOException e) {
                        System.err.println("IP定向消息发送失败: " + session.getId());
                    }
                });
    }

    // 复用现有的IP提取逻辑
    private String getIpFromSession(WebSocketSession session) {
        try {
            InetSocketAddress socketAddress = (InetSocketAddress) session.getRemoteAddress();
            InetAddress address = socketAddress.getAddress();

            // 处理IPv6环回地址的情况
            if (address instanceof Inet6Address && address.isLoopbackAddress()) {
                return "127.0.0.1"; // 显式返回IPv4环回地址
            }

            // 包含IPv6和IPv4的转换逻辑（与getAllClientIps方法一致）
            if (address instanceof Inet6Address) {
                byte[] bytes = address.getAddress();
                if (bytes.length == 16 &&
                        bytes[0] == 0 && bytes[1] == 0 &&
                        bytes[2] == 0 && bytes[3] == 0 &&
                        bytes[4] == 0 && bytes[5] == 0 &&
                        bytes[6] == 0 && bytes[7] == 0 &&
                        bytes[8] == 0 && bytes[9] == 0 &&
                        bytes[10] == (byte) 0xff && bytes[11] == (byte) 0xff) {

                    // 提取真实的IPv4地址
                    return String.format("%d.%d.%d.%d",
                            bytes[12] & 0xff,
                            bytes[13] & 0xff,
                            bytes[14] & 0xff,
                            bytes[15] & 0xff);
                }
            }
            return address.getHostAddress();
        } catch (Exception e) {
            return "未知IP";
        }
    }
}
